import seaborn as sns
from matplotlib import pyplot as plt
import os
from sklearn.metrics import r2_score
# sns.set_style("darkgrid")
fontsize=20

def plot(x, ys, ynames: list, xlabel: str, ylabel: str, save_path: str = 'results'):
    """
    Args:
        x: array([int, ]) or list[], x轴的坐标
        ys: [[int, ..], [int, ..]],  二级列表，每个元素列表都是一条线
        ynames: list[str, ] ,  二级列表，和ys对应，为图例的名称
        xlabel: str, x轴名称
        ylabel: str, y轴名称
        save_path: str, 保存路径
    """
    # plt.rcParams['font.sans-serif']=['Arial Unicode Ms']
    sns.set(context='notebook', style='darkgrid', font_scale=2)
    fig = plt.figure(figsize=(12, 10))
    for y, name in zip(ys, ynames):
        plt.plot(x, y, label=name)

    plt.legend(loc=1)
    plt.xticks(x[::len(x)//5], fontsize=13)
    plt.ylabel(ylabel)
    plt.savefig(os.path.join(save_path, 'plot.jpg'), dpi=600)
    

def dfplot(df, save_path):
    sns.set(context='notebook', style='darkgrid', font_scale=2)
    os.makedirs(save_path, exist_ok=True)
    fig = plt.figure(figsize=(12, 10))
    sns.lineplot(data=df)

    plt.savefig('results/plot.png', dpi=400)


def plot2(x, ys, n, xlabel, ylabel, save_path: str = 'results', filename=None) -> None:
    """
    绘制散点图
    """
    # print(len(x))
    print(f"n_{filename}={len(x)}")
    sns.set(context='notebook', style='darkgrid', font_scale=1.2)
    plt.figure(figsize=(13, 13))
    os.makedirs(save_path, exist_ok=True)
    # print(len(set(x)))
    # print(len(set(ys)))
    plt.scatter(x, ys, s=34)
    midx, midy = sum(x)/len(x), sum(ys)/len(x)
    
    minx, miny, maxx, maxy = min(x), min(ys), max(x), max(ys)
    mmin, mmax = min(minx, miny), max(maxx, maxy)
    plt.plot([mmin, mmax], [mmin, mmax], label="1:1 line", color='red', linewidth=2)
    plt.plot([minx, maxx], [miny, maxy], label="Trendline")
    r2 = r2_score(ys, x)
    # plt.scatter(midx, midy, label=f"$R^2$={round(r2, 6)}", alpha=0)
    plt.xlabel(xlabel+"(m)")
    plt.ylabel(ylabel+"(m)")
    plt.xlim([mmin-100, mmax+100])
    plt.ylim([mmin-100, mmax+100])
    plt.scatter(midx, midy, label=f"n={n}", alpha=0)
    plt.legend(loc='best')
    filename = filename if filename else 'scatter'
    plt.savefig(os.path.join(save_path, f'{filename}.jpg'), dpi=400)